package com.example.videopractice.mediaex

import android.media.*
import android.os.Environment
import android.util.Log
import com.example.base.util.ALog
import com.example.videopractice.mediaex.H264Encode.MuxerData.Type.Audio
import com.example.videopractice.mediaex.H264Encode.MuxerData.Type.Video
import kotlinx.coroutines.Dispatchers.IO
import kotlinx.coroutines.GlobalScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import java.io.File
import java.nio.ByteBuffer
import java.util.concurrent.LinkedBlockingQueue

class H264Encode {

    // 视频相关参数
    private val width = 640
    private val height = 480
    private val frameRate = 30

    private var videoCodec: MediaCodec? = null
    //private val yuv420Queue = ArrayBlockingQueue<ByteArray>(10)
    private val yuv420Queue = LinkedBlockingQueue<ByteArray>()

    // 音频相关参数
    private val sampleRateInHz: Int = 44100
    private val channelConfig: Int = AudioFormat.CHANNEL_CONFIGURATION_MONO
    private val encodingBitRate: Int = AudioFormat.ENCODING_PCM_16BIT
    private var audioCodec: MediaCodec? = null
    private var audioRecord: AudioRecord? = null

    private var isRunning = false
    private var outputPath =
        Environment.getExternalStorageDirectory().absolutePath + File.separator + "xiao.mp4"
    private val muxerDateQueue = LinkedBlockingQueue<MuxerData>()
    private var mediaMuxer: MediaMuxer? = null
    private var runningTime = 0L

    @Volatile
    private var isVideoTrackAdd = false
    @Volatile
    private var isAudioTrackAdd = false

    fun startEncoder() {
        runningTime = System.currentTimeMillis()
        isRunning = true
        muxerDateQueue.clear()
        isVideoTrackAdd = false
        isAudioTrackAdd = false
        GlobalScope.launch(IO) {
            mediaMuxer = MediaMuxer(outputPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
            launch { catch({ onVideoEncoder() }, "onVideoEncoder") }
            launch { catch({ onAudioEncoder() }, "onAudioEncoder") }
            launch { catch({ onMuxer()}, "onMuxer") }
        }
    }

    fun stopEncoder() {
        ALog.e("xiao","录制时长: ${(System.currentTimeMillis() - runningTime) / 1000}")
        isRunning = false
    }

    fun putVideoData(buffer: ByteArray) {
//        if (yuv420Queue.size >= 10) {
//            yuv420Queue.poll()
//        }
        yuv420Queue.put(buffer)
    }

    private var pts = 0L
    private fun onVideoEncoder() {
        val muxer = mediaMuxer ?: return

        var videoTrack = -1
        val mediaFormat =
            MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height)
        mediaFormat.setInteger(
            MediaFormat.KEY_COLOR_FORMAT,
            MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar
        )
        //COLOR_FormatYUV420SemiPlanar
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, width * height * 5)
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 30)
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1)

        videoCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC)
        val videoCodec = videoCodec ?: throw NullPointerException("videoEncoder 未初始化")
        videoCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        videoCodec.start()

        var generateIndex = 0L
        while (isRunning) {
            val input = yuv420Queue.take()
            val yuv420sp = ByteArray(width * height * 3 / 2)
            // 必须要转格式，否则录制的内容播放出来为绿屏
            NV21ToNV12(input, yuv420sp, width, height)

            try {
                //获取输入流队列
                val inputBufferIndex = videoCodec.dequeueInputBuffer(0) //间隔多久 0为立即返回
                val inputBuffer = videoCodec.getInputBuffer(inputBufferIndex)
                if (inputBuffer != null) {
                    pts = computePresentationTime(generateIndex)
                    inputBuffer.put(yuv420sp)
                    videoCodec.queueInputBuffer(inputBufferIndex, 0, yuv420sp.size, pts, 0)
                    generateIndex += 1
                }

                //获取输出流队列
                val bufferInfo = MediaCodec.BufferInfo()
                var outputBufferIndex = videoCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                    videoTrack = muxer.addTrack(videoCodec.outputFormat)
                    isVideoTrackAdd = true
                }
                while (outputBufferIndex >= 0) {
                    if (videoTrack >= 0) {
                        videoCodec.getOutputBuffer(outputBufferIndex)?.let {
                            if (isMuxerStart()) muxerDateQueue.put(MuxerData(Video, videoTrack, it ,bufferInfo))
                        }
                        videoCodec.releaseOutputBuffer(outputBufferIndex, false)
                    }
                    outputBufferIndex = videoCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                }
            } catch (e: Exception) {
                ALog.e("xiao", "onVideoEncoder出错: ${e.message}")
                e.printStackTrace()
            }
        }

        try {
            videoCodec.stop()
            videoCodec.release()
            this.videoCodec = null
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    private fun onAudioEncoder() {
        val muxer = mediaMuxer ?: return
        var audioTrack = -1

        val audioFormat =
            MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_AAC, sampleRateInHz, channelConfig)
        audioFormat.setInteger(MediaFormat.KEY_AAC_PROFILE, MediaCodecInfo.CodecProfileLevel.AACObjectLC)
        audioFormat.setInteger(MediaFormat.KEY_BIT_RATE, 64 * 1000)
        audioCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_AUDIO_AAC)
        audioCodec?.configure(audioFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        audioCodec?.start()

        val recordBufSize = AudioRecord.getMinBufferSize(sampleRateInHz, channelConfig, encodingBitRate)
        audioRecord = AudioRecord(MediaRecorder.AudioSource.MIC, sampleRateInHz, channelConfig, encodingBitRate, recordBufSize)
        audioRecord?.startRecording()

        val audioCodec = audioCodec ?: throw NullPointerException()
        val audioRecord = audioRecord ?: throw NullPointerException()

        val bufferInfo = MediaCodec.BufferInfo()
        val buffer = ByteArray(recordBufSize)
        while (isRunning) {
            val readBytes = audioRecord.read(buffer, 0, recordBufSize)
            try {
                val inputBufferIndex = audioCodec.dequeueInputBuffer(TIMEOUT_USEC)
                if (inputBufferIndex < 0) return
                val inputBuffer = audioCodec.getInputBuffer(inputBufferIndex)
                if (inputBuffer != null) {
                    inputBuffer.put(buffer)
                    if (readBytes <= 0) {
                        ALog.e("xiao", "send BUFFER_FLAG_END_OF_STREAM")
                        audioCodec.queueInputBuffer(inputBufferIndex, 0, 0, pts, MediaCodec.BUFFER_FLAG_END_OF_STREAM)
                    } else {
                        audioCodec.queueInputBuffer(inputBufferIndex, 0, readBytes, pts, 0)
                    }
                }

                /*获取解码后的数据*/
                var outputBufferIndex = audioCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                    audioTrack = muxer.addTrack(audioCodec.outputFormat)
                    if (audioTrack >= 0) {
                        isAudioTrackAdd = true
                    }
                }

                while (outputBufferIndex >= 0) {
                    if (audioTrack >= 0) {
                        val outBuffer = audioCodec.getOutputBuffer(outputBufferIndex)
                        if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG != 0) {
                            bufferInfo.size = 0
                        }
                        if (bufferInfo.size != 0 && outBuffer != null) {
                            if (isMuxerStart()) muxerDateQueue.put(MuxerData(Audio, audioTrack, outBuffer, bufferInfo))
                        }
                        audioCodec.releaseOutputBuffer(outputBufferIndex, false)
                    }
                    outputBufferIndex = audioCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                }

            } catch (e: Exception) {
                ALog.e("xiao", "解码音频(Audio)数据失败: ${e.message}")
                e.printStackTrace()
            }
        }

        try {
            audioRecord.stop()
            audioRecord.release()
            this.audioCodec = null
        } catch (e: Exception) {
            ALog.e("xiao", e.message)
            e.printStackTrace()
        }

        try {
            audioCodec.stop()
            audioCodec.release()
            this.audioCodec = null
        } catch (e: Exception) {
            ALog.e("xiao", e.message)
            e.printStackTrace()
        }
    }

    private suspend fun onMuxer() {
        while (!isMuxerStart()) { delay(100) }
        val muxer = mediaMuxer ?: throw NullPointerException()
        muxer.start()

        ALog.e("xiao","开始muxer")
        while (isRunning || muxerDateQueue.size > 0) {
            muxerDateQueue.take().apply {
                Log.e("xiao","type: $type")
                muxer.writeSampleData(trackIndex, buffer, bufferInfo)
            }
        }

        try {
            muxer.stop()
            muxer.release()
        } catch (e: Exception) {
            ALog.e("xiao", e.message)
            e.printStackTrace()
        }
        mediaMuxer = null
        ALog.e("xiao","muxer结束")
    }

    private fun NV21ToNV12(nv21: ByteArray, nv12: ByteArray, width: Int, height: Int) {
        val frameSize = width * height
        System.arraycopy(nv21, 0, nv12, 0, frameSize)
        for (i in 0 until frameSize) {
            nv12[i] = nv21[i]
        }
        for (j in 0 until (frameSize / 2) step 2) {
            nv12[frameSize + j - 1] = nv21[j + frameSize]
        }
        for (j in 0 until (frameSize / 2) step 2) {
            nv12[frameSize + j] = nv21[j + frameSize - 1]
        }
    }

    /**
     * 根据帧数生成时间戳
     */
    private fun computePresentationTime(frameIndex: Long): Long {
        return 132 + frameIndex * 1000000 / frameRate
    }

    private fun isMuxerStart() = isAudioTrackAdd && isVideoTrackAdd

    private suspend fun catch(action: (suspend () -> Unit), key: String) {
        try {
            action.invoke()
        } catch (e: Exception) {
            ALog.e("xiao", "$key ${e.message}")
            e.printStackTrace()
        }
    }

    data class MuxerData(
        val type: Type,
        val trackIndex: Int,
        val buffer: ByteBuffer,
        val bufferInfo: MediaCodec.BufferInfo
    ) {
        enum class Type {
            Video,
            Audio
        }
    }

    companion object {
        private const val TIMEOUT_USEC = 12000L
    }
}